import os
import rasterio
from rasterio.windows import Window
import numpy as np

# 定义输入和输出路径
input_path = '../../data/picture/origin_tif/EdgeCut1.tif'  # 替换为你的输入文件路径
output_dir = '../../data/picture/output_tif/train/label'  # 替换为你的输出目录路径

# 确保输出目录存在
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 打开 TIFF 文件
with rasterio.open(input_path) as src:
    # 获取图像的宽度和高度
    width, height = src.width, src.height
    # 定义窗口大小
    win_size = 256  # 256x256 像素

    # 计算需要裁剪的窗口数量
    for i in range(0, height, win_size):
        for j in range(0, width, win_size):
            # 定义窗口的左上角坐标
            ul_x = j
            ul_y = height - i - win_size

            # 创建窗口对象
            window = Window(ul_x, ul_y, win_size, win_size)

            # 读取窗口内的数据
            data = src.read(window=window)

            # 创建输出文件名
            output_filename = f"patch_{i}_{j}.tif"
            output_path = os.path.join(output_dir, output_filename)

            # 写入数据到新的 TIFF 文件
            with rasterio.open(
                    output_path,
                    'w',
                    driver='GTiff',
                    height=win_size,
                    width=win_size,
                    count=src.count,
                    dtype=data.dtype,
                    crs=src.crs,
                    transform=rasterio.windows.transform(window, src.transform)
            ) as dest:
                dest.write(data)

print("裁剪完成。")